#!/usr/bin/python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import json
import struct
from collections import namedtuple
from typing import List, NamedTuple, Optional, Tuple


_EDID_STRUCT_FORMAT = (
    '>'     # big-endian
    '8s'    # constant header (8 bytes)
    'H'     # manufacturer id (2 bytes)
    'H'     # product id (2 bytes)
    'I'     # serial number (4 bytes)
    'B'     # manufactoring week (1 byte)
    'B'     # manufactoring year (1 byte)
    'B'     # edid version (1 byte)
    'B'     # edid revision (1 byte)
    'B'     # video input type (1 byte)
    'B'     # horizontal size in cm (1 byte)
    'B'     # vertical size in cm (1 byte)
    'B'     # display gamma (1 byte)
    'B'     # supported features (1 byte)
    '10s'   # color characteristics (10 bytes)
    'H'     # supported timings (2 bytes)
    'B'     # reserved timing (1 byte)
    '16s'   # EDID supported timings (16 bytes)
    '18s'   # detailed timing block 1 (18 bytes)
    '18s'   # detailed timing block 2 (18 bytes)
    '18s'   # detailed timing block 3 (18 bytes)
    '18s'   # detailed timing block 4 (18 bytes)
    'B'     # extension flag (1 byte)
    'B'     # checksum (1 byte)
)


_RAW_EDID = namedtuple(
    '_RAW_EDID',
    (
        'header',
        'manu_id',
        'prod_id',
        'serial_no',
        'manu_week',
        'manu_year',
        'edid_version',
        'edid_revision',
        'input_type',
        'width',
        'height',
        'gamma',
        'features',
        'color',
        'timings_supported',
        'timings_reserved',
        'timings_edid',
        'timing_1',
        'timing_2',
        'timing_3',
        'timing_4',
        'extension',
        'checksum'
    )
)


_TIMINGS = {
    0:  (1280, 1024, 75.),
    1:  (1024,  768, 75.),
    2:  (1024,  768, 72.),
    3:  (1024,  768, 60.),
    4:  (1024,  768, 87.),
    5:  ( 832,  624, 75.),
    6:  ( 800,  600, 75.),
    7:  ( 800,  600, 70.),
    8:  ( 800,  600, 60.),
    9:  ( 800,  600, 56.),
    10: ( 640,  480, 75.),
    11: ( 640,  480, 72.),
    12: ( 640,  480, 67.),
    13: ( 640,  480, 60.),
    14: ( 720,  400, 88.),
    15: ( 720,  400, 70.),
}


_ASPECT_RATIOS = {
    0b00: (16, 10),
    0b01: ( 4,  3),
    0b10: ( 5,  4),
    0b11: (16,  9),
}


# Parsed EDID object
class EDID(NamedTuple):
    manufacturer_id: int
    product_id: int
    year: int
    week: int
    edid_version: str
    type: str
    width: float
    height: float
    gamma: float
    dpms_standby: bool
    dpms_suspend: bool
    dpms_activeoff: bool
    resolutions: List[Tuple[int, int, float]]
    name: Optional[str]
    serial: Optional[str]

    def __repr__(self) -> str:
        return f"EDID({self.manufacturer} {self.name}, s/n {self.serial})"

    def __str__(self) -> str:
        return json.dumps(self._asdict())


class EdidUtil:

    def parse_edid(raw):
        # crop additional edid data
        if len(raw) > 128:
            raw = raw[:128]

        if sum(raw) % 256 != 0:
            raise ValueError('Checksum mismatch')

        if raw[:8] != b'\x00\xff\xff\xff\xff\xff\xff\x00':
            raise ValueError('Invalid EDID header')

        unpacked = struct.unpack(_EDID_STRUCT_FORMAT, raw)
        raw_edid = _RAW_EDID(*unpacked)

        resolutions = []
        for i in range(16):
            bit = raw_edid.timings_supported & i
            if bit:
                resolutions.append(_TIMINGS[16-i])

        for i in range(8):
            try:
                bytes_data = raw_edid.timings_edid[2*i:2*i+2]
                if bytes_data == b'\x01\x01':
                    continue
                byte1, byte2 = bytes_data
                x_res = 8*(int(byte1)+31)
                aspect_ratio = _ASPECT_RATIOS[(byte2 >> 6) & 0b11]
                y_res = int(x_res * aspect_ratio[1]/aspect_ratio[0])
                rate = (int(byte2) & 0b00111111) + 60.0
                resolutions.append((x_res, y_res, rate))
            except BaseException:
                pass

        name: Optional[str] = None
        serial: Optional[str] = None

        for timing_bytes in (raw_edid.timing_1, raw_edid.timing_2, raw_edid.timing_3, raw_edid.timing_4):
            # 'other' descriptor
            if timing_bytes[0:2] == b'\x00\x00':
                timing_type = timing_bytes[3]
                if timing_type in (0xFF, 0xFE, 0xFC):
                    buffer = timing_bytes[5:]
                    buffer = buffer.partition(b'\x0a')[0]
                    text = buffer.decode('cp437')
                    if timing_type == 0xFF:
                        serial = text
                    elif timing_type == 0xFC:
                        name = text

        return EDID(
            manufacturer_id=raw_edid.manu_id,
            product_id=raw_edid.prod_id,
            year=raw_edid.manu_year + 1990,
            week=raw_edid.manu_week,
            edid_version='{:d}.{:d}'.format(raw_edid.edid_version, raw_edid.edid_revision),
            type='digital' if (raw_edid.input_type & 0xFF) else 'analog',
            width=float(raw_edid.width),
            height=float(raw_edid.height),
            gamma=(raw_edid.gamma + 100) / 100,
            dpms_standby=bool(raw_edid.features & 0xFF),
            dpms_suspend=bool(raw_edid.features & 0x7F),
            dpms_activeoff=bool(raw_edid.features & 0x3F),
            resolutions=resolutions,
            name=name,
            serial=serial if serial else raw_edid.serial_no
        )
